/**
 * 
 */
package edu.hitsz.cs.nlp.application;

import java.util.HashMap;
import java.util.Vector;

import edu.hitsz.cs.ml.onlinelearning.AveragedPerceptron;
import edu.hitsz.cs.nlp.struct.ConllFile;
import edu.hitsz.cs.nlp.struct.ConllSentence;
import edu.hitsz.cs.nlp.util.FeatureTemplate;
import edu.hitsz.cs.util.ArgumentsParser;


/**
 * @author tm
 * Averaged Perceptron Algorithm for Sequence Labeling Problems
 */
public class SequenceAP {

	public SequenceAP(){
	}
	
	/**
	 * To find the Viterbi path using parameters without Averaged perceptron 
	 * @param ap AveragedPerceptron
	 * @param file Conll00File
	 * @param feaStrategy Strategy of feature selection: 0 represents from template, 1 from hand-crafted.
	 * @param i The ith sentence
	 * @return
	 */
	public void viterbi(AveragedPerceptron ap,ConllFile file, FeatureTemplate newTemplate, int i){
		ConllSentence sentence=file.totalSentence.get(i);
		int sentenceLen=sentence.sentenceLength;
		sentence.predictSignal.clear();
		for(int j=0;j<sentenceLen;j++){
			Vector<String> oneFeature=new Vector<String>();
			sentence.extractTokenFeaturesFromTmpt(oneFeature, newTemplate,j);			
			double weight=0;
			String tmpResultSignal=file.resultSignal.get(0);
			//iterative every possible tag
			for(int k=0;k<file.resultSignal.size();k++){
				double tmpWeight=ap.compSingleWeights(oneFeature, file.resultSignal.get(k));
				if(tmpWeight>weight){
					weight=tmpWeight;
					tmpResultSignal=file.resultSignal.get(k);
				}
			}
			sentence.word.get(j).add(tmpResultSignal);
			if (sentence.predictSignal.size()==0){
				Vector<String> newVector=new Vector<String> ();
				sentence.predictSignal.add(newVector);
			}
			sentence.predictSignal.get(0).add(tmpResultSignal);
		}
		//file.printFeatures(sentence.predicateSignal);
		//file.printFeatures(sentence.senResultSignal);
	}
	
	/**
	 * To find the Viterbi path using parameters with Averaged perceptron 
	 * @param ap AveragedPerceptron
	 * @param file Conll00File
	 * @param feaStrategy Strategy of feature selection: 0 represents from template, 1 from hand-crafted.
	 * @param i The ith sentence
	 * @return
	 */
	public void viterbiAP(AveragedPerceptron ap,ConllFile file, FeatureTemplate newTemplate, int i){
		ConllSentence sentence=file.totalSentence.get(i);
		int sentenceLen=sentence.sentenceLength;
		sentence.predictSignal.clear();
		for(int j=0;j<sentenceLen;j++){
			Vector<String> oneFeature=new Vector<String>();
			
			//extract features from template, or from hand-crafted codes
			sentence.extractTokenFeaturesFromTmpt(oneFeature, newTemplate, j);
						
			double weight=0;
			String tmpResultSignal=file.resultSignal.get(0);
			//iterative every possible tag
			for(int k=0;k<file.resultSignal.size();k++){
				double tmpWeight=ap.compAllWeights(oneFeature, file.resultSignal.get(k));
				if(tmpWeight>weight){
					weight=tmpWeight;
					tmpResultSignal=file.resultSignal.get(k);
				}
			}
			sentence.word.get(j).add(tmpResultSignal);
			if (sentence.predictSignal.size()==0){
				Vector<String> newVector=new Vector<String> ();
				sentence.predictSignal.add(newVector);
			}
			sentence.predictSignal.get(0).add(tmpResultSignal);
		}
		//file.printFeatures(sentence.predicateSignal);
		//file.printFeatures(sentence.senResultSignal);
	}
	
	/**
	 * To find the Viterbi path using parameters with Averaged perceptron 
	 * @param ap AveragedPerceptron
	 * @param file Conll00File
	 * @param newTemplate feature template
	 * @param feaStrategy Strategy of feature selection: 0 represents from template, 1 from hand-crafted.
	 * @param ithSentence ith sentence
	 * @param kb number of beams
	 * @return
	 */
	public void kBeam(AveragedPerceptron ap,ConllFile file, FeatureTemplate newTemplate, int ithSentence, int kb){
		ConllSentence sentence=file.totalSentence.get(ithSentence);	
		//temporary vectors
		Vector<Vector<String>> allSignal=new Vector<Vector<String>>();
		Vector<Double> allWeight=new Vector<Double>();
		int sentenceLen=sentence.sentenceLength;
		//an initial value
		sentence.predictValue.add(0.0);
		sentence.predictSignal.clear();
		Vector<String> newSignal=new Vector<String>();
		sentence.predictSignal.add(newSignal);
		//iterate every token in the sentence
		for(int i=0;i<sentenceLen;i++){
			int minK=Math.min(kb, sentence.predictSignal.size());
			//iterate every previous result in previous sequence
			allSignal.clear();
			allWeight.clear();
			for(int j=0;j<minK;j++){
				//record temporary value
				double weight=sentence.predictValue.get(j);
				Vector<String> oneFeature=new Vector<String>();				
				sentence.extractTokenFeaturesFromTmpt(oneFeature, newTemplate, i, j);				
				
				//iterative every possible tag
				for(int m=0;m<file.resultSignal.size();m++){					
					double tmpWeight=weight+ap.compAllWeights(oneFeature, file.resultSignal.get(m));
					//
					if(allWeight.size()==0){
						Vector<String> newVector=new Vector<String>();
						newVector.addAll(sentence.predictSignal.get(j));
						newVector.add(file.resultSignal.get(m));
						allSignal.add(newVector);
						allWeight.add(tmpWeight);						
					}
					else{
						int minNext=Math.min(kb, allWeight.size());
						int n=0;
						for(;n<minNext;n++){
							if(allWeight.get(n)>=tmpWeight)
								continue;
							else
								break;
						}
						if(n<kb){
							Vector<String> newVector=new Vector<String>();
							newVector.addAll(sentence.predictSignal.get(j));
							newVector.add(file.resultSignal.get(m));
							allSignal.add(n,newVector);
							allWeight.add(n,tmpWeight);
						}						
					}
				}
			}
			int maxNext=Math.min(kb, allWeight.size());
			sentence.predictSignal.clear();
			sentence.predictValue.clear();
			sentence.predictSignal.addAll(allSignal.subList(0, maxNext));
			sentence.predictValue.addAll(allWeight.subList(0, maxNext));			
		}
		/*
		int maxNum=Math.min(kb, allWeight.size());
		for(int i=0;i<maxNum;i++)
			sentence.printFeatures(sentence.predictSignal.get(i));
		*/
		
	}

	/**
	 * update the parameter for Averaged Perceptron
	 * @param ap
	 * @param file
	 * @param feaStrategy
	 * @param ithSentence 
	 */
	public void update(AveragedPerceptron ap,ConllFile file, FeatureTemplate newTemplate,int ithSentence){
		ConllSentence sentence=file.totalSentence.get(ithSentence);
		int sentenceLen=sentence.sentenceLength;
		HashMap<String,Integer> senFeature=new HashMap<String,Integer>();		
		for(int j=0;j<sentenceLen;j++){
			Vector<String> oneFeature=new Vector<String>();
			//extract features from template
			sentence.extractTokenFeaturesFromTmpt(oneFeature, newTemplate, j);
			//update the parameter of every parameter
			for(int o=0;o<oneFeature.size();o++){
				String one=oneFeature.get(o)+"+"+sentence.senResultSignal.get(j);
				if(senFeature.containsKey(one)){
					int num=senFeature.get(one);
					senFeature.put(one, num+1);
				}
				else
					senFeature.put(one, 1);
				
				one=oneFeature.get(o)+"+"+sentence.predictSignal.get(0).get(j);
				if(senFeature.containsKey(one)){
					int num=senFeature.get(one);
					senFeature.put(one, num-1);
				}
				else
					senFeature.put(one, -1);
			}			
		}
		/*
		if(sentence.senResultSignal!=sentence.predictSignal){
			ap.update(trueFeature, trueFeatureNum,predicateFeature,predicateFeatureNum);
		}	
		*/
		if((!sentence.senResultSignal.equals(sentence.predictSignal))||(ithSentence+1==file.sentenceNumber))
			ap.update(senFeature,file.sentenceNumber,ithSentence);
	}
	
	
	
	/**
	 * Train the classifier
	 * @param iterNum the number of iteration
	 * @param obj String of objection, such as "pos", "chunk"
	 */
	public void train(String trainName, String templateName, String signalName, String modelName, int iterNum){
		ConllFile trainFile=new ConllFile();		
		FeatureTemplate newTemplate=new FeatureTemplate();
		newTemplate.readFromFile(templateName);	
		trainFile.readTrainFile(trainName, 0);	
		trainFile.generateResultSignal();
		trainFile.storeResultSignal(signalName);
		AveragedPerceptron ap=new AveragedPerceptron();
		for(int i=1;i<=iterNum;i++){
			System.out.println("\nThis is the "+Integer.toString(i)+"th iteration...");
			int sentenceNum=trainFile.sentenceNumber;
			System.out.print("Sentence ");
			for(int j=0;j<sentenceNum;j++){	
				if(j%100==0){
					System.out.print(Integer.toString(j)+"...");
					if(j%1000==0)
						System.out.println();
				}
				viterbi(ap,trainFile,newTemplate,j);				
				update(ap,trainFile,newTemplate,j);				
			}		
			System.out.println(Integer.toString(sentenceNum));
			String singleFeatureName=modelName+"-"+Integer.toString(i);
			ap.storeWeights(singleFeatureName,i,sentenceNum);
		}		
		System.out.println("\nTraining process is done!");		
	}
	
	
	/**
	 * Predict results using the classifier in all iteration
	 * @param testName
	 * @param templateName
	 * @param signalName
	 * @param modelName
	 * @param resultName
	 * @param iterNum
	 */
	public void predictAll(String testName, String templateName,String signalName,String modelName, String resultName, int iterNum){
		ConllFile testFile=new ConllFile();
		FeatureTemplate newTemplate=new FeatureTemplate();
		newTemplate.readFromFile(templateName);	
		testFile.readTestFile(testName, 0);		
		testFile.readResultSignal(signalName);
		testFile.printVector(testFile.resultSignal);
		int i=1;
		for(;i<=iterNum;i++)
		{
			System.out.println("This is "+Integer.toString(i)+"th iteration");
			String featureName=modelName+"-"+i;		
			String finalResultName=resultName+"-"+i;
			AveragedPerceptron ap=new AveragedPerceptron();
			ap.readWeights(featureName);
			int sentenceNum=testFile.sentenceNumber;
			System.out.print("\nCompute Viterbi Path for Sentence ");		
			int j=0;
			for(;j<sentenceNum;j++){				
				if(j%100==0){
					System.out.print(Integer.toString(j)+"...");
					if(j%1000==0)
						System.out.println();
				}
				viterbiAP(ap,testFile,newTemplate,j);				
			}
			System.out.println("Total Sentence "+Integer.toString(j)+" have been processed.");
			testFile.storePredictFile(finalResultName);						
		}
	}

	/**
	 * Predict results using the classifier
	 * @param testName
	 * @param templateName
	 * @param signalName
	 * @param modelName
	 * @param resultName
	 * @param kNum
	 */
	public void predict(String testName, String templateName,String signalName,String modelName, String resultName, int kNum){
		ConllFile testFile=new ConllFile();
		FeatureTemplate newTemplate=new FeatureTemplate();
		newTemplate.readFromFile(templateName);	
		testFile.readTestFile(testName, 0);		
		testFile.readResultSignal(signalName);
		testFile.printVector(testFile.resultSignal);
		
		String featureName=modelName;		
		String finalResultName=resultName;
		AveragedPerceptron ap=new AveragedPerceptron();
		ap.readWeights(featureName);
		int sentenceNum=testFile.sentenceNumber;
		System.out.print("\nCompute Viterbi Path for Sentence ");		
		int j=0;
		for(;j<sentenceNum;j++){				
			if(j%100==0){
				System.out.print(Integer.toString(j)+"...");
				if(j%1000==0)
					System.out.println();
			}
			if(kNum==1)
				viterbiAP(ap,testFile,newTemplate,j);
			else if(kNum>1)
				kBeam(ap,testFile,newTemplate,j,kNum);
		}
		System.out.println("Total Sentence "+Integer.toString(j)+" have been processed.");
		if(kNum==1)
			testFile.storePredictFile(finalResultName);
		else if(kNum>1)
			for(int k=0;k<kNum;k++){
				testFile.storePredictFile(finalResultName, k);
			}				
	}
	
	
	public void predict(String testName, String templateName,String signalName,String modelName, String resultName){
		predict(testName,templateName,signalName,modelName, resultName,1);
	}
	
	

	
	static void usage(){
		System.out.println("Usage:");
		System.out.println("  Help: -h");
		System.out.println("  Training: -t trainFile -e templateFile -s signalFile -m modelFile -i iterativeNumber");
		System.out.println("    You need prepare the trainFile and templateFile, and set the iterationNumber");
		System.out.println("  Predicting iteration: -p testFile -e templateFile -s signalFile -m modelFile -r resultName -i iterativeNumber");
		System.out.println("  Predicting kbest: -p testFile -e templateFile -s signalFile -m modelFile -r resultName -k kBestNum");
		System.out.println("  Predicting one: -p testFile -e templateFile -s signalFile -m modelFile -r resultName");
		System.exit(1);
	}
	
	
	/**
	 * 
	 * @param args
	 */
	static void oneExample(String[] args){
		SequenceAP newW=new SequenceAP();
				
		String trainName="";
		String templateName="";
		String signalName="";
		String modelName="";
		String testName="";
		String resultName="";
		int iterateTimes=1;		
		int kNum=1;
		
		String shortArgs="ht:e:s:m:p:r:i:k:";
		ArgumentsParser newParser=new ArgumentsParser();
		newParser.parseCmdLine(args, shortArgs);
		
		if(newParser.containsArgument("t"))
			trainName=newParser.getArgument("t");
		if(newParser.containsArgument("e"))
			templateName=newParser.getArgument("e");
		if(newParser.containsArgument("s"))
			signalName=newParser.getArgument("s");
		if(newParser.containsArgument("m"))
			modelName=newParser.getArgument("m");
		if(newParser.containsArgument("p"))
			testName=newParser.getArgument("p");
		if(newParser.containsArgument("r"))
			resultName=newParser.getArgument("r");
		if(newParser.containsArgument("i"))
			iterateTimes=Integer.parseInt(newParser.getArgument("i"));
		if(newParser.containsArgument("k"))
			kNum=Integer.parseInt(newParser.getArgument("k"));
		
		//
		if(newParser.containsOption("h"))
			usage();
		else if(newParser.containsArgument("t")&&newParser.containsArgument("e")&&newParser.containsArgument("s")&&newParser.containsArgument("m")&&newParser.containsArgument("i"))
				newW.train(trainName,templateName,signalName,modelName,iterateTimes);
		else if(newParser.containsArgument("p")&&newParser.containsArgument("e")&&newParser.containsArgument("s")&&newParser.containsArgument("m")&&newParser.containsArgument("r")){
			if(newParser.containsArgument("i")){
				if(!newParser.containsArgument("k"))
					newW.predictAll(testName,templateName,signalName,modelName,resultName,iterateTimes);
				else 
					usage();
			}				
			else
				if(!newParser.containsArgument("k"))
					newW.predict(testName,templateName,signalName,modelName,resultName);
				else
					newW.predict(testName,templateName,signalName,modelName,resultName,kNum);
		}
		else 
			usage();		
		//newW.predict(20,1,obj);
	}
	
	
	
	public static void main(String[] args){
		//System.out.println(false^false);
		oneExample(args);		
		
	}

}
